from bot.enterprise_wechat_gpt.util.llm_agent_aio import *
from bot.enterprise_wechat_gpt.util.merge_conversation_util import *
from bot.enterprise_wechat_gpt.test_case.extract_user_situation_case import *
from bot.enterprise_wechat_gpt.util.id_distributor import distribute


class TestAgent(LLMAgent):
    prompt_template = """
    角色：{salesperson_role}
    任务：你要判断用户的经济情况，比如有钱还是没钱
    
    请回复：
    """

    def __init__(self, conversation_history):  # conversation_history是整合同一角色的连续多行对话 为一行后，得到的对话历史list
        self.conversation_str = '\n'.join([list(d.keys())[0] + ':' + list(d.values())[0] for d in conversation_history])
        self.prompt = self.prompt_template.format(salesperson_role=SYSTEM, conversation_history=self.conversation_str)
        super().__init__(self.prompt)


class UserSituation(LLMAgent):
    """抽取用户情境Agent"""
    user_question_v0 = """
        作为{salesperson_role}，在微信对话中，你的任务是根据与用户的对话来识别和判断用户当前的借款情境。用户情境可以包括以下几种：
        1. 用户身份：学生、军人、上班族、企业主、个体户、无业。
        2. 用户年龄：18岁以下、18-55岁、56-64岁、65以上。
        3. 借款用途：装修、医疗看病、生意周转、养殖、囤货、日常消费、结婚、买车、开店、租房、教育培训等。
        4. 借款金额：本次借贷的金额。
        5. 借款时长：借款一个月、借款两年等。
        6. 何时借款：判断用户是否未来有借贷需求或者周期性借款需求，1个月之后用、两个月之后用、月初再借、月底使用、9月份使用、每个月都借。
        7. 逾期情况：当上下文中提到"逾期"两个字时，请判断逾期情况。比如"无逾期"、"逾期90天"。如果用户只是咨询，不用返回该情境。
        8. 咨询产品：当前用户正在咨询的产品，包括金条、白条、分分卡、信用卡、企业主贷、借钱、黑卡。
        9. 当前提额状态：判断用户去金条提额入口"提升额度"的状态，还未提额、提额成功、提额失败。
        10. 当前借款状态：判断用户金条借款的状态，还未借款、审核中、到账途中、借款成功、借款失败、无法借款、借不了。
        11. 当前优惠券状态：判断用户金条优惠券的状态，想要优惠券、已获得优惠券、已使用优惠券。还款券不要识别为这个情境。
        12. 当前账户状态：判断用户金条账户的状态，无可用额度、额度低、额度正常、账户不能用了、利率高。只要当前可用额度小于借款金额就是额度低。

        注意事项:
        1.输出格式：key-value格式，key为情境，value为情境的描述。情境和描述之间用"："分割，同一个情境的多个描述之间用"、"分割，多个情境之间用"｜"分割。
        2.对于未提及的情境，无需例举出来。如果没有识别到任何情境，请直接返回"无情境"。
        3.请根据与用户的对话记录，回答用户当前的情境和情境的描述，不需要回答用户的具体问题。
        4."额度是否提取"表示"当前借款状态"，而不是"当前提额状态"。

        对话纪录如下：
        ===
        {conversation_history}
        ===
        识别的情境：
    """
    user_question_v1 = """
        作为{salesperson_role}，在微信对话中，你的任务是根据与用户的对话来识别和判断用户当前的借款情境。用户情境可以包括以下几种：
    
        1. 当前提额的状态：还没有提额、提额成功、提额失败。
        2. 当前借款状态：还没有借款、审核中、到账途中、借款成功、借款失败、无法借款、借不了。
        3. 当前优惠券状态：想要优惠券、已获得优惠券、已使用优惠券。还款券不要识别为这个情境。
        4. 当前账户状态：无可用额度、额度低、额度正常、账户不能用了、利率高。只要当前可用额度小于借款金额就是额度低。
        5. 借款用途：装修、医疗看病、生意周转、养殖、囤货、日常消费、结婚、买车、开店、租房、教育培训等。
        6. 借款金额：本次借贷的金额。
        7. 借款时长：借款一个月、借款两年。
        8. 何时借款：判断用户是否未来有借贷需求或者周期性借款需求，1个月之后用、两个月之后用、月初再借、月底使用、9月份使用、每个月都借。
        9. 逾期情况：当上下文中提到"逾期"两个字时，请判断逾期情况。比如"无逾期"、"逾期90天"。如果用户只是咨询，不用返回该情境。
        10. 产品：当前用户正在咨询的产品，包括金条、白条、分分卡、信用卡、企业主贷、借钱、黑卡。
        11. 用户身份：学生、军人、上班族、企业主、个体户、无业。
        12. 用户年龄：18岁以下、18-55岁、56-64岁、65以上。
        
        注意事项:
        1.输出格式：key-value格式，key为情境，value为情境的描述。情境和描述之间用"："分割，同一个情境的多个描述之间用"、"分割，多个情境之间用"｜"分割。
        2.对于未提及的情境，无需例举出来。如果没有识别到任何情境，请直接返回"无情境"。
        3.请根据与用户的对话记录，回答用户当前的情境和情境的描述，不需要回答用户的具体问题。
        
        示例开始：
        用户：额度怎么没有了，利息也太高了
        回答：当前账户状态：无可用额度，利率高
        用户：1
        回答：无情境
        用户：想借5w
        回答：当前借款状态：还没有借款
        用户：我的分分卡没额度了
        回答：产品：分分卡｜当前账户状态：无可用额度
        示例结束
        
        对话纪录如下：
        ===
        {conversation_history}
        ===
        识别的情境：
    """

    def __init__(self, conversation_history):
        self.conversation_str = transfer_conversation_list_to_str(conversation_history)
        self.prompt = self.user_question_v0.format(salesperson_role=SYSTEM, conversation_history=self.conversation_str)
        super().__init__(self.prompt)


async def get_user_situation(deal_conversation_history):
    user_id = local_data.user_id
    user_situation = UserSituation(deal_conversation_history)
    if distribute(user_id, 2):
        res_user_situation = await user_situation.achat_llm(agent_name='用户情境')
    else:
        res_user_situation = await user_situation.achat_llm(agent_name='用户情境', model_version='wechat_financial_7b_1218')
    logging.info(f"userId:[{user_id}] 用户情境: {res_user_situation}")
    try:
        list = [i for i in res_user_situation.split("｜") if not ("无情境" in i or "无法确定" in i or i == '无' or "未提及" in i)]
        res_user_situation = "｜".join(list)
    except Exception as e:
        logging.error(f"获取用户情境失败: 返回：{res_user_situation}, 错误：{e}")
    return res_user_situation


async def get_test(conversation_history):
    test_agent = TestAgent(conversation_history)
    res_test = await test_agent.chat_with_gpt(temperature=0.9, agent_name='Agent1.3', model='gpt-3.5-turbo')
    logging.info(f"=====> 1.3 Test_Result: {res_test} <=====")
    return res_test

if __name__ == '__main__':
    history = parse_history(dialogue_text)
    deal_conversation_history = merge_role_conversation(history)
    loops = asyncio.get_event_loop()
    res_user_situation = loops.run_until_complete(get_user_situation(deal_conversation_history))
    print(res_user_situation)




